## Code: ui.py
```python
## ui.py

import pygame
from typing import Dict, List
from constants import COLORS, TILE_SIZE, MARGIN, DIRECTIONS


class UI:
    """
    Class to handle the user interface for the 2048 game. It is responsible for rendering the board,
    displaying the score, handling user input, and showing game over screens.
    """

    def __init__(self, board_size: int, difficulty: str = "NORMAL"):
        """
        Initialize the UI class with the board size and difficulty level.

        Args:
        board_size (int): The size of the board.
        difficulty (str): The difficulty level of the game. Default is "NORMAL".
        """
        self._board_size = board_size
        self._difficulty = difficulty
        self._screen = pygame.display.set_mode((
            (TILE_SIZE + MARGIN) * board_size + MARGIN, 
            (TILE_SIZE + MARGIN) * board_size + MARGIN
        ))
        pygame.display.set_caption(f"2048 Game - Difficulty: {self._difficulty}")

    def display_score(self, score: int) -> None:
        """
        Display the current score on the screen.

        Args:
        score (int): The current score to display.
        """
        font = pygame.font.Font(None, 36)
        score_text = font.render(f"Score: {score}", True, (0, 0, 0))
        score_rect = score_text.get_rect(topleft=(MARGIN, MARGIN))
        self._screen.blit(score_text, score_rect)

    def display_game_over(self) -> None:
        """
        Display the game over message on the screen.
        """
        font = pygame.font.Font(None, 72)
        game_over_text = font.render("GAME OVER", True, (255, 0, 0))
        game_over_rect = game_over_text.get_rect(center=(
            (TILE_SIZE + MARGIN) * self._board_size / 2,
            (TILE_SIZE + MARGIN) * self._board_size / 2
        ))
        self._screen.blit(game_over_text, game_over_rect)

    def display_restart_button(self) -> None:
        """
        Display the restart button on the screen.
        """
        font = pygame.font.Font(None, 36)
        restart_text = font.render("Restart", True, (0, 0, 0))
        restart_rect = restart_text.get_rect(topleft=(
            (TILE_SIZE + MARGIN) * self._board_size - 100,
            MARGIN
        ))
        pygame.draw.rect(self._screen, (255, 255, 255), (restart_rect.left - 5, restart_rect.top - 5, 110, 40))
        self._screen.blit(restart_text, restart_rect)

    def render_board(self, board: 'Board') -> None:
        """
        Render the game board on the screen.

        Args:
        board (Board): The game board to render.
        """
        self._screen.fill((0, 0, 0))
        for i in range(self._board_size):
            for j in range(self._board_size):
                tile_value = board._tiles[i][j]
                color = COLORS[tile_value]["background"]
                text_color = COLORS[tile_value]["foreground"]
                pygame.draw.rect(self._screen, color, (
                    (TILE_SIZE + MARGIN) * j + MARGIN,
                    (TILE_SIZE + MARGIN) * i + MARGIN,
                    TILE_SIZE,
                    TILE_SIZE
                ))
                if tile_value != 0:
                    font = pygame.font.Font(None, 48)
                    text = font.render(str(tile_value), True, text_color)
                    text_rect = text.get_rect(center=(
                        (TILE_SIZE + MARGIN) * j + MARGIN + TILE_SIZE / 2,
                        (TILE_SIZE + MARGIN) * i + MARGIN + TILE_SIZE / 2
                    ))
                    self._screen.blit(text, text_rect)
        pygame.display.flip()

    def handle_input(self, event: pygame.event.Event, game: 'Game') -> None:
        """
        Handle user input events such as key presses and restart button clicks.

        Args:
        event (pygame.event.Event): The event to handle.
        game (Game): The game instance to interact with.
        """
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_UP:
                game.move_tiles(DIRECTIONS["UP"])
            elif event.key == pygame.K_DOWN:
                game.move_tiles(DIRECTIONS["DOWN"])
            elif event.key == pygame.K_LEFT:
                game.move_tiles(DIRECTIONS["LEFT"])
            elif event.key == pygame.K_RIGHT:
                game.move_tiles(DIRECTIONS["RIGHT"])
        elif event.type == pygame.MOUSEBUTTONDOWN:
            mouse_pos = pygame.mouse.get_pos()
            if (self._board_size * (TILE_SIZE + MARGIN) - 100 + MARGIN) > mouse_pos[0] > (